import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from work_sequence_plan_utils.work_sequence_plan_transform import cal_order_sequence
from matplotlib.patches import Patch

def getMachineTypeColor(row):
    return c_dict[row['machine_type']]

if __name__ == '__main__':
    # 订单信息
    order_info = [
        ['A', '001', 3, 20],
        ['A', '002', 3, 50],
        ['B', '003', 4, 50],
        ['B', '004', 4, 20],
        ['B', '005', 4, 30],
        ['A', '006', 3, 50],
        ['B', '007', 4, 50],
        ['B', '008', 4, 50],
        ['A', '009', 3, 40],
        ['B', '010', 4, 50],
    ]
    order_num = len(order_info)
    transform_cost = [
        {'name': 'A-A', 'cost': 1},
        {'name': 'B-B', 'cost': 1},
        {'name': 'A-B', 'cost': 3},
        {'name': 'B-A', 'cost': 2},
    ]
    # 结果信息
    order_sequence_list = [3,8,5,0,1,4,2,6,9,7]
    order_sequence_list_info = cal_order_sequence(order_sequence_list, order_info, transform_cost)
    # print(order_sequence_list_info)

    order_sequence_df = pd.DataFrame(order_sequence_list_info)
    order_sequence_df.columns = ['machine_type', 'order_no', 'product_time', 'due_date',
                                 'start_time', 'end_time', 'tranform_cost']
    # 颜色信息
    c_dict = {'A': '#3475D0', 'B': '#E69646'}
    legend_elements = [Patch(facecolor=c_dict[i], label=i) for i in c_dict]
    print (legend_elements)
    order_sequence_df['color'] = order_sequence_df.apply(getMachineTypeColor, axis=1)
    # print(order_sequence_df)
    # 排序
    sort_order_sequence_df = order_sequence_df.sort_values(by=['order_no'], ascending=[True])
    # print(sort_order_sequence_df)

    # 作图
    fig, ax = plt.subplots(1, figsize=(10, 4))
    ax.barh(sort_order_sequence_df.order_no, sort_order_sequence_df.product_time,
            left=sort_order_sequence_df.start_time, color=sort_order_sequence_df.color)
    ax.barh(sort_order_sequence_df.order_no, sort_order_sequence_df.tranform_cost,
            left=sort_order_sequence_df.start_time - sort_order_sequence_df.tranform_cost,
            color='#34D05C',alpha=0.5)
    ax.barh(sort_order_sequence_df.order_no, 0.2,
            left=sort_order_sequence_df.due_date,
            color='#E64646')
    plt.legend(handles=legend_elements)


    plt.show()

    pass
